import gradio as gr
from PIL import Image
import numpy as np
from sklearn.cluster import KMeans

def extract_palette(image, progress=gr.Progress()):
    """
    从PIL图像中提取5种主色，并返回带样式的HTML结果
    """
    if image is None:
        return "<p style='color: #999; text-align: center;'>请上传一张图片</p>"

    progress(0.1, desc="正在转换图像格式...")
    image = image.convert("RGB")
    
    progress(0.3, desc="正在缩小图像以加速处理...")
    image = image.resize((100, 100), Image.LANCZOS)
    
    progress(0.5, desc="正在提取像素数据...")
    pixels = np.array(image).reshape(-1, 3)
    
    progress(0.7, desc="正在进行K-Means聚类...")
    kmeans = KMeans(n_clusters=5, random_state=42, n_init=10)
    kmeans.fit(pixels)
    
    progress(0.9, desc="正在生成调色板...")
    colors = kmeans.cluster_centers_.astype(int)
    hex_colors = [f"#{r:02x}{g:02x}{b:02x}" for r, g, b in colors]
    
    progress(1.0, desc="完成！")

    # 构建带内联样式的HTML（放大显示区域）
    html_output = '''
    <div style="
        background-color: #1e1e1e;
        color: white;
        padding: 20px;
        border-radius: 12px;
        font-family: 'Segoe UI', sans-serif;
        max-width: 600px;
        margin: 0 auto;
        box-shadow: 0 4px 12px rgba(0,0,0,0.3);
    ">
        <h3 style="margin-top: 0; text-align: center; color: #f0f0f0;">🎨 主色调调色板</h3>
        
        <!-- 颜色块 -->
        <div style="
            display: flex;
            justify-content: center;
            gap: 8px;
            margin-bottom: 20px;
            flex-wrap: wrap;
        ">
    '''

    for color in hex_colors:
        html_output += f'''
        <div style="
            width: 90px;
            height: 60px;
            background-color: {color};
            border-radius: 6px;
            box-shadow: inset 0 0 0 1px rgba(255,255,255,0.1);
        "></div>
        '''

    html_output += '''
        </div>

        <!-- 十六进制值 -->
        <div style="
            display: flex;
            justify-content: center;
            flex-wrap: wrap;
            gap: 12px;
            background-color: #2a2a2a;
            padding: 14px;
            border-radius: 8px;
            font-family: monospace;
            font-size: 15px;
        ">
    '''

    for color in hex_colors:
        html_output += f'<span style="color: {color}; font-weight: bold;">{color}</span>'

    html_output += '''
        </div>
    </div>
    '''

    return html_output

# 创建 Gradio 界面
iface = gr.Interface(
    fn=extract_palette,
    inputs=gr.Image(type="pil", label="🖼️ 上传图片"),
    outputs=gr.HTML(label=""),
    title="🎨 图像主色调调色板提取器",
    description="上传任意图片，自动提取其中5种主要颜色及其十六进制值",
    examples=[
        ["data/example1.jpg"],
        ["data/example2.png"]
    ],  # 可选：添加示例图片路径（确保 data/ 目录存在）
    cache_examples=False,
)

if __name__ == "__main__":
    iface.launch()